/**
* \brief 
* \author pengcheng (pengcheng@yslrpch@126.com)
* \date 2020-05-20
* \attention CopyrightÃ‚Â©ADC Technology(tianjin)Co.Ltd
* \attention Refer to COPYRIGHT.txt for complete terms of copyright notice
*/
#include <iostream>
#include "cmw_app/sgm_app.h"
#include <yaml-cpp/yaml.h>
#include <opencv2/calib3d/calib3d.hpp>
#include <opencv2/opencv.hpp>
#include <chrono>
#include "common/project_root.h"

namespace adc
{

StereoSgmApp::StereoSgmApp(const std::string &config_file, const float hz)  : Worker(hz)
{
    InitConfig(config_file);
    point_color_callback_ = NULL;
    point_cloud_callback_ = NULL;
    poind_and_color_callback_ = NULL;
    point_cloud_mat_callback_ = NULL;
}

void StereoSgmApp::SetPointCloudCallback(PointCloudCallback callback)
{
    point_cloud_callback_ = callback;
}

void StereoSgmApp::SetPointColorCallback(PointColorCallback callback)
{
    point_color_callback_ = callback;
}


void StereoSgmApp::SetPointAdnColorCallback(PointAndColorCallback callback)
{
    poind_and_color_callback_ = callback;
}

void StereoSgmApp::SetPointCloudMatCallback(PointCloudMatCallback callback)
{
    point_cloud_mat_callback_ = callback;
}
void StereoSgmApp::InitConfig(const std::string &config_file)
{
    try
    {
        adc::RootPath::Ptr  root_path_ptr = adc::RootPath::GetInstance();
        YAML::Node node = YAML::LoadFile(config_file);
        disp_size_ = node["disp_size"].as<int>(64);
        output_depth_ = node["output_depth"].as<int>(16);
        subpixel_ = node["subpixel"].as<bool>(true);
        P1_ = node["p1"].as<int>(10);
        P2_ = node["p2"].as<int>(120);
        input_type_ = node["input_type"].as<int>(8);
        uniqueness_ = node["uniqueness"].as<float>(0.95f);
        int image_height = node["image_height"].as<int>();
        int image_width = node["image_width"].as<int>();
        std::string left_camera_params = adc::RootPath::GetAbsolutePath(node["left_camera_yaml"].as<std::string>());
        std::string right_camera_params = adc::RootPath::GetAbsolutePath(node["right_camera_yaml"].as<std::string>());

        float camera_height = node["camera_height"].as<float>(1.65);
        float camera_tilt = node["camera_tilt"].as<float>(0.0);
        float  camera_yaw = node["camera_yaw"].as<float>(0.0);
        float camera_pitch = node["camera_pitch"].as<float>(0.0);
        top_row_ = node["top_row"].as<int>();
        max_depth_ = node["max_depth"].as<float>();
        point_cloud_mat_ = cv::Mat(image_height, image_width, CV_32FC3, cv::Scalar::all(-1000));
        input_bytes_ = input_type_ * image_height * image_width / 8;
        output_bytes_ = output_depth_ * image_height * image_width / 8;
        dispartity_ = cv::Mat(image_height, image_width, CV_16S);
        d_left_ptr_ = std::make_shared<device_buffer>(input_bytes_);
        d_right_ptr_ = std::make_shared<device_buffer>(input_bytes_);
        d_dispartity_ptr_ = std::make_shared<device_buffer>(output_bytes_);
        adc::StereoSGM::Parameters param(P1_, P2_, uniqueness_, subpixel_);
        sgm_ptr_ = std::make_shared<adc::StereoSGM>(image_width, image_height, disp_size_, input_type_, output_depth_,
                adc::EXECUTE_INOUT_CUDA2CUDA, param);

        cv::Mat K_l, D_l, R_l, P_l, K_r, D_r, R_r, P_r;
        cv::Size size_l, size_r;
        StereoSgmApp::ReadParams(left_camera_params, K_l, D_l, P_l, R_l, size_l);
        StereoSgmApp::ReadParams(right_camera_params, K_r, D_r, P_r, R_r, size_r);

        cv::initUndistortRectifyMap(K_l, D_l, R_l, P_l, size_l, CV_32FC1, map_l_x, map_l_y);
        cv::initUndistortRectifyMap(K_r, D_r, R_r, P_r, size_r, CV_32FC1, map_r_x, map_r_y);

        camera_ptr_ = std::make_shared<CameraParameters>();
        camera_ptr_->fu = P_l.at<double>(0,0);
        camera_ptr_->fv = P_l.at<double>(1,1);
        camera_ptr_->u0 = P_l.at<double>(0,2);
        camera_ptr_->v0 = P_l.at<double>(1,2);
        camera_ptr_->baseline = -P_r.at<double>(0,3) / P_r.at<double>(0,0);
        camera_ptr_->height = camera_height;
        camera_ptr_->tilt = camera_tilt;
        camera_ptr_->yaw = camera_yaw;
        camera_ptr_->pitch = camera_pitch;
         tf_ptr_ = std::make_shared<CoordinateTransform>(*camera_ptr_);
    }
    catch(const std::exception& e)
    {
        std::cerr << e.what() << '\n';
    }
}

void StereoSgmApp::ReadParams(const std::string& yaml_file, cv::Mat& K, cv::Mat& D, cv::Mat& P, cv::Mat& R, cv::Size& size)
{
     try{
        YAML::Node node = YAML::LoadFile(yaml_file);
        int height = node["image_height"].as<int>();
        int width = node["image_width"].as<int>();
        size = cv::Size2i(width, height);
        std::vector<double> K_v= node["camera_matrix"]["data"].as<std::vector<double>>();
        std::vector<double> P_v = node["projection_matrix"]["data"].as<std::vector<double>>();
        std::vector<double> D_v = node["distortion_coefficients"]["data"].as<std::vector<double>>();
        std::vector<double> R_v = node["rectification_matrix"]["data"].as<std::vector<double>>();

        cv::Mat _K(K_v);
        cv::Mat _P(P_v);
        cv::Mat _D(D_v);
        cv::Mat _R(R_v);
        K = _K.reshape(3,3).clone();
        P = _P.reshape(4,3).clone();
        D = _D.reshape(14, 1).clone();
        R = _R.reshape(3,3).clone();
        }
        catch(const std::exception& e)
        {
            std::cerr << e.what() << '\n'<<std::endl;;
            // ROS_ERROR(e.what());
            exit(1);
        }
}

void StereoSgmApp::CacheImages(const cv::Mat &left_img, const cv::Mat &right_img, std::chrono::system_clock::time_point &time)
{
    std::lock_guard<std::mutex> mutex_cache_image(mutex_image_);
    
    image_cache_.push(std::make_pair(left_img, right_img));
    time_cache_.push(time);
    while (image_cache_.size() > kQueueSize)
    {
        image_cache_.pop();
        time_cache_.pop();
    }
}

bool StereoSgmApp::ReceiveImage(void)
{
    std::lock_guard<std::mutex> mutex_cache_image(mutex_image_);
    if(image_cache_.empty())
    {
        return false;
    }
    left_image_ = image_cache_.front().first.clone();
    right_image_ = image_cache_.front().second.clone();
    current_time_ = time_cache_.front();
    image_cache_.pop();
    time_cache_.pop();
    return true;
}

void StereoSgmApp::StereoMatch()
{
    cv::Mat left_rect, right_rect;
    cv::remap(left_image_, left_image_, map_l_x, map_l_y, cv::INTER_LINEAR);
    cv::remap(right_image_, right_image_, map_r_x, map_r_y, cv::INTER_LINEAR);
    cv::cvtColor(left_image_, left_rect, cv::COLOR_BGR2GRAY);
    cv::cvtColor(right_image_, right_rect, cv::COLOR_BGR2GRAY);
    // cv::imshow("left_rect", left_rect);
    // cv::imshow("right_rect", right_rect);
    cudaMemcpy(d_left_ptr_->data, left_rect.data, input_bytes_, cudaMemcpyHostToDevice);
    cudaMemcpy(d_right_ptr_->data, right_rect.data, input_bytes_, cudaMemcpyHostToDevice);
    sgm_ptr_->execute(d_left_ptr_->data, d_right_ptr_->data, d_dispartity_ptr_->data);
    cudaMemcpy(dispartity_.data, d_dispartity_ptr_->data, output_bytes_, cudaMemcpyDeviceToHost);
    dispartity_.convertTo(dispartity_32f_, CV_32F, subpixel_ ? 1. / adc::StereoSGM::SUBPIXEL_SCALE : 1.0);
}

void StereoSgmApp::ReprojectPointsTo3D()
{
    CV_Assert(dispartity_32f_.type() == CV_32F);
    point_cloud_.clear();
    point_color_.clear();
    point_cloud_.reserve(dispartity_32f_.rows * dispartity_32f_.cols);
    point_color_.reserve(dispartity_32f_.rows * dispartity_32f_.cols);

    for (int y = top_row_; y < dispartity_32f_.rows ; y++)
	{
        float* cpm_ptr = point_cloud_mat_.ptr<float>(y);
		for (int x = 0; x < dispartity_32f_.cols; x++)
		{
			const float d = dispartity_32f_.at<float>(y, x);
			if (d > 0)
			{
				cv::Point3f temp = tf_ptr_->imageToWorld(cv::Point(x, y), d);
				if(temp.z < max_depth_)
                {
                    cv::Vec3b color = left_image_.at<cv::Vec3b>(y,x);
					point_cloud_.push_back(temp);
                    point_color_.push_back(color);
                    cpm_ptr[x*3] = temp.x;
                    cpm_ptr[x*3 + 1] = temp.y;
                    cpm_ptr[x*3 + 2] = temp.z;

                }
                else
                {
                    cpm_ptr[x*3] = -1000.0;
                    cpm_ptr[x*3 + 1] = -1000.0;
                    cpm_ptr[x*3 + 2] = -1000.0;
                }
			}
            else
            {
                cpm_ptr[x*3] = -1000.0;
                cpm_ptr[x*3 + 1] = -1000.0;
                cpm_ptr[x*3 + 2] = -1000.0;
            }
		}
	}

}

void StereoSgmApp::DoWork()
{
    static int count = 0;
    if(!ReceiveImage())
    {
        if(count > 50)
        {
            std::cout<<"can not get images"<<std::endl; 
        }
        count ++;
        return;
    }
    count = 0;
   // const auto t1 = std::chrono::system_clock::now();
    StereoMatch();
    //const auto t2 = std::chrono::system_clock::now();
   // const auto duration = std::chrono::duration_cast<std::chrono::microseconds>(t2 - t1).count();
    ReprojectPointsTo3D();
    //const auto fps = 1e6 / duration;
    //std::cout<<"match fps: "<<fps<<std::endl;
    //std::cout<<"generate points: "<<point_cloud_.size()<<std::endl;
    if(NULL != point_color_callback_)
    {
        point_color_callback_(point_color_, current_time_);
    }
    if(NULL != point_cloud_callback_)
    {
        point_cloud_callback_(point_cloud_, current_time_);
    }
    if(NULL != poind_and_color_callback_)
    {
        poind_and_color_callback_(point_cloud_,  point_color_, current_time_);
    }
    if(NULL != point_cloud_mat_callback_)
    {
        // cv::imshow("left", );
        point_cloud_mat_callback_(point_cloud_mat_, current_time_);
    }
}


}
